from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr
from matplotlib.widgets import Button
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import PolynomialFeatures
from scipy.optimize import curve_fit
from scipy.fft import fft2,ifft2
from scipy.fft import fftn, ifftn
from scipy.interpolate import Rbf
import scipy.signal
from scipy.ndimage import gaussian_filter
from ripser import ripser
from persim import plot_diagrams
from scipy.spatial.distance import pdist, squareform

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal


def preprocess(trjs, max_length):
    processed_trjs = []
    for i in range(trjs.shape[0]):
        trj = trjs[i]
        if trj.shape[0] < max_length:
            last_element = trj[-1]
            processed_trj = np.concatenate([trj, np.repeat(last_element[np.newaxis,:], max_length - trj.shape[0], axis=0)], axis=0)
            processed_trjs.append(processed_trj)
        else:
            processed_trjs.append(trj)
    processed_trjs = np.array(processed_trjs)
    return processed_trjs
    
@jax.jit
def build_radiance_field(p):
    img = jnp.zeros((21, 21))
    top = 10
    bottom = 0
    effective_radius = 21*1.414
    px, py = p[0], p[1]

    x_coord_map = jnp.arange(img.shape[0])
    x_coord_map = jnp.repeat(x_coord_map[:, jnp.newaxis], img.shape[1], axis=1)

    y_coord_map = jnp.arange(img.shape[1])
    y_coord_map = jnp.repeat(y_coord_map[jnp.newaxis, :], img.shape[0], axis=0)
    
    dist = jnp.sqrt((x_coord_map - px)**2 + (y_coord_map - py)**2)
    img = jnp.where(dist < effective_radius,
                    (top - bottom) * (effective_radius - dist) / effective_radius + bottom,
                    img)
    return img

@jax.jit
def get_max_radiance_field(imgs):
    img = jnp.max(imgs, axis=0)
    return img

@jax.jit
def build_ridge(A):

    # 将 A 整体平移到(10,10)


    # 使用 jnp.map 来并行计算每个辐射场图像
    imgs = jax.vmap(build_radiance_field)(A)
    img = get_max_radiance_field(imgs)
    return img

build_ridge_vmap = jax.vmap(build_ridge)


random_group = np.load("./logs/random_group.npy")
# contat_group = np.load("./logs/contat_group.npy")

print(random_group.shape)

concat_group = []
for i in range(10000):
    progress_bar(i, 10000)
    
    # 从random_group中随机挑选4个元素，把他们拼接到一起
    idx1 = random.randint(0, random_group.shape[0]-1)
    idx2 = random.randint(0, random_group.shape[0]-1)
    idx3 = random.randint(0, random_group.shape[0]-1)
    idx4 = random.randint(0, random_group.shape[0]-1)

    trj1 = random_group[idx1]
    trj2 = random_group[idx2]
    trj3 = random_group[idx3]
    trj4 = random_group[idx4]

    end_trj1 = trj1[-1]
    trj2 = trj2-trj2[0]+end_trj1
    end_trj2 = trj2[-1]
    trj3 = trj3-trj3[0]+end_trj2
    end_trj3 = trj3[-1]
    trj4 = trj4-trj4[0]+end_trj3

    new_trj = np.concatenate([trj1, trj2, trj3, trj4], axis=0)

    # 生成一个 (5,5) 范围内的随机偏移
    dx = random.randint(0, 5)
    dy = random.randint(0, 5)
    new_trj = new_trj + np.array([dx, dy])

    concat_group.append(new_trj)

concat_group = np.array(concat_group)
print(concat_group.shape)

ridge_images = build_ridge_vmap(concat_group)

for u in range(10):
    # 随机选择一个ridge_images显示出来
    idx = random.randint(0, ridge_images.shape[0]-1)
    plt.imshow(ridge_images[idx], cmap='viridis', interpolation='nearest')
    plt.show()

# 将 RIs 从 (100, 21, 21) 转换成 (100, 441)
ridge_images = ridge_images.reshape(ridge_images.shape[0], ridge_images.shape[1] * ridge_images.shape[2])
print("RIs.shape: ", ridge_images.shape)

# 对 ridge_images 进行 PCA
pca = PCA()
pca.fit(ridge_images)
ridge_images_pca = pca.transform(ridge_images)

imgs_reshaped_pca_components = pca.components_

# 将 ridge_images_pca_grouped 绘制到3D视图中
fig = plt.figure()
ax1 = fig.add_subplot(111)

view_dim0 = 0

def button_clicked1(event):

    global view_dim0

    if view_dim0 <= imgs_reshaped_pca_components.shape[0]-2:
        view_dim0 += 1
    # 在这里清空原有的图像
    ax1.clear()
    
    # 显示 ridge_pca_components[view_dim0-2] 的 21*21 图像
    ax1.imshow(imgs_reshaped_pca_components[view_dim0].reshape(21,21))
    
    ax1.set_title('view_dim0: ' + str(view_dim0))
    # 刷新图像
    plt.draw()


def button_clicked2(event):
    global view_dim0
    if view_dim0 > 0:
        view_dim0 -= 1
    # 在这里清空原有的图像
    ax1.clear()
    
    # 显示 ridge_pca_components[view_dim0-2] 的 21*21 图像
    ax1.imshow(imgs_reshaped_pca_components[view_dim0].reshape(21,21))
    
    ax1.set_title('view_dim0: ' + str(view_dim0))
    # 刷新图像
    plt.draw()

# 创建按钮的位置和大小
button_ax1 = fig.add_axes([0.6, 0.1, 0.1, 0.1])
button_ax2 = fig.add_axes([0.3, 0.1, 0.1, 0.1])
# 创建按钮对象，并设置按钮的文本和位置
button1 = Button(button_ax1, 'dim_forward')
button2 = Button(button_ax2, 'dim_backward')
# 绑定按钮的点击事件和响应函数
button1.on_clicked(button_clicked1)
button2.on_clicked(button_clicked2)

# 显示 ridge_pca_components[view_dim0-2] 的 21*21 图像
ax1.imshow(imgs_reshaped_pca_components[view_dim0].reshape(21,21))

ax1.set_title('view_dim0: ' + str(view_dim0))
# 刷新图像
plt.show()

# 绘制 PCA 的 variance ratio
fig = plt.figure()
ax = fig.add_subplot(111)
ax.bar(range(len(pca.explained_variance_ratio_)), pca.explained_variance_ratio_)
ax.set_title('explained variance ratio')
plt.show()

